package zy.learn.demo.structuredstreaming.watermark

import java.sql.Timestamp

import org.apache.spark.SparkConf
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.streaming.{OutputMode, Trigger}

/**
 * 输入自定义的时间戳
 * 水印的作用是抛弃过期的窗口，而不是数据，晚于水印的数据依然可以参与不晚于水印的窗口里的计算
 * Complete Output模式下，水印没有意义
 */
object Watermark1 {
  def main(args: Array[String]): Unit = {
    val sparkConf = new SparkConf().set("spark.sql.shuffle.partitions", "3")
    val spark = SparkSession.builder()
      .master("local[2]")
      .config(sparkConf)
      .appName("Watermark")
      .getOrCreate()

    import spark.implicits._

    val lines = spark.readStream
      .format("socket") // 设置数据源
      .option("host", "co7-203")
      .option("port", 9999)
      .load
    /* 输入的数据
    * 2020-10-14 10:55:00,dog     # wartermark: 10:55:00 - 2 min = 10:53:00
    * 2020-10-14 11:00:00,dog     # wartermark: 11:00:00 - 2 min = 10:58:00
    * 2020-10-14 10:55:00,dog     # wartermark: 10:55:00 - 2 min = 10:53:00 < 10:58:00 => 10:58:00
    * 输出的结果
+------------------------------------------+-----+-----+
|window                                    |word |count|
+------------------------------------------+-----+-----+
|[2020-10-14 10:50:00, 2020-10-14 11:00:00]|dog |1    |   Batch 0
|[2020-10-14 10:46:00, 2020-10-14 10:56:00]|dog |1    |
|[2020-10-14 10:48:00, 2020-10-14 10:58:00]|dog |1    |
|[2020-10-14 10:52:00, 2020-10-14 11:02:00]|dog |1    |
|[2020-10-14 10:54:00, 2020-10-14 11:04:00]|dog |1    |
+------------------------------------------+-----+-----+
|[2020-10-14 10:56:00, 2020-10-14 11:06:00]|dog |1    |   Batch 1
|[2020-10-14 10:52:00, 2020-10-14 11:02:00]|dog |2    |
|[2020-10-14 10:54:00, 2020-10-14 11:04:00]|dog |2    |
|[2020-10-14 10:58:00, 2020-10-14 11:08:00]|dog |1    |
|[2020-10-14 11:00:00, 2020-10-14 11:10:00]|dog |1    |
+------------------------------------------+----+-----+
|[2020-10-14 10:50:00, 2020-10-14 11:00:00]|dog |2    |   Batch 2
|[2020-10-14 10:52:00, 2020-10-14 11:02:00]|dog |3    |   [2020-10-14 10:46:00, 2020-10-14 10:56:00]和[2020-10-14 10:48:00, 2020-10-14 10:58:00]
|[2020-10-14 10:54:00, 2020-10-14 11:04:00]|dog |3    |   窗口因水印为 10:58:00 而被抛弃掉
+------------------------------------------+----+-----+
    * */
    val wordsDF = lines.as[String].map(line => {
      val split = line.split(",")
      (Timestamp.valueOf(split(0)), split(1))
    }).toDF("ts", "word")

    import org.apache.spark.sql.functions._

    val wordCounts = wordsDF
      .withWatermark("ts", "2 minutes")
      .groupBy(
        // 调用 window 函数, 返回的是一个 Column 参数 1: df 中表示时间戳的列 参数 2: 窗口长度 参数 3: 滑动步长
        window($"ts", "10 minutes", "2 minutes"),
        $"word"
      ).count() // 计数

    val query = wordCounts.writeStream
      .format("console")
      .outputMode(OutputMode.Update())
      .trigger(Trigger.ProcessingTime(2000))
      .option("truncate", "false")
      .start()

    query.awaitTermination()
  }
}
